package frontend.visitor;

import exceptions.VisitInnerException;
import frontend.visitor.symboltbl.SymbolTbl;
import ir.types.datatypes.DataType;
import ir.types.datatypes.IntegerType;
import ir.types.datatypes.PointerType;
import ir.types.datatypes.VoidType;
import ir.values.*;
import ir.values.users.GlobalVariable;
import ir.values.users.constants.Constant;
import ir.values.users.constants.IntegerConst;
import ir.values.users.instructions.Instruction;
import ir.values.users.instructions.jump.JumpInst;
import ir.values.users.instructions.mem.AllocaInst;
import ir.values.users.instructions.mem.GetElementPtrInst;
import ir.values.users.instructions.mem.StoreInst;
import utils.BBList;

import java.util.ArrayList;
import java.util.List;
import java.util.Stack;

public class IRBuildFactory {
    private final SymbolTbl symbolTbl;
    private final ModuleIr moduleIr = new ModuleIr();
    private FunctionIr curFunctionIr;
    private BBList<BasicBlock>.BBNode curBasicBlockNode;
    private BBList<BasicBlock>.BBNode curBasicBlockNodeEnd;

    public IRBuildFactory(SymbolTbl symbolTbl) {
        this.symbolTbl = symbolTbl;

        moduleIr.addFunction(GET_INT);
        moduleIr.addFunction(PUT_INT);
        moduleIr.addFunction(PUT_CHAR);
        moduleIr.addFunction(PUT_STR);
    }
    private static ArrayList<Argument> buildArgumentList(Argument... arguments) {
        return new ArrayList<>(List.of(arguments));
    }

    public static FunctionIr GET_INT = new FunctionIr(
            "getint", true, IntegerType.I32, buildArgumentList());
    public static FunctionIr PUT_INT = new FunctionIr(
            "putint", true, VoidType.VOID, buildArgumentList(
            new Argument("a", IntegerType.I32)
    )
    );
    public static FunctionIr PUT_CHAR = new FunctionIr(
            "putch", true, VoidType.VOID, buildArgumentList(
            new Argument("a", IntegerType.I32)
    )
    );
    public static FunctionIr PUT_STR = new FunctionIr(
            "putstr", true, VoidType.VOID, buildArgumentList(
            new Argument("str", new PointerType(IntegerType.I8))
    )
    );

    public FunctionIr buildFunction(String name, DataType returnedType, ArrayList<Argument> arguments) {
        this.endOfBasicBlock = false;
        if (!symbolTbl.inRootScope()) throw new VisitInnerException();

        // 构建函数对象
        FunctionIr functionIr = new FunctionIr(name, false, returnedType, arguments);

        // 进入函数首个基本块
        this.moduleIr.addFunction(functionIr);
        this.curFunctionIr = functionIr;
        switchCurBasicBlock(functionIr.getBasicBlocks().getHeadNode());

        return functionIr;
    }

    public void buildFunctionIrPointerArguments() {
        for (Argument argument: curFunctionIr.getArguments()) {
            AllocaInst allocaInst = buildAllocaInst(argument.getOriginalName(), argument.getDataType());
            StoreInst storeInst = new StoreInst(argument, allocaInst);
            buildInst(storeInst);
            argument.setPointerValue(allocaInst);
        }
    }

    public DataType getCurFunctionIrReturnType() {
        return curFunctionIr.getFunctionType().getReturnType();
    }

    public GlobalVariable buildGlobalVariable(String name, Constant constant, boolean isConst) {
        if (!symbolTbl.inRootScope()) throw new VisitInnerException();
        GlobalVariable globalVariable = new GlobalVariable(name, constant, isConst);
        moduleIr.addGlobalVariable(globalVariable);
        return globalVariable;
    }

    private boolean endOfBasicBlock = false;

    public <T extends Instruction> T buildInst(T instruction) {
        if (endOfBasicBlock) {
            switchCurBasicBlock(createBasicBlock());
            endOfBasicBlock = false;
        }
        curBasicBlockNode.get().addInstruction(instruction);
        if (instruction instanceof JumpInst) {
            endOfBasicBlock = true;
        }
        return instruction;
    }

    public AllocaInst buildAllocaInst(String ident, DataType baseType) {
        AllocaInst allocaInst = new AllocaInst(ident, new PointerType(baseType), symbolTbl.peekScope());
        return buildInst(allocaInst);
    }

    public Value tryBuildGetElementPtrInst(Value pointerValue, ArrayList<Value> indexValues) {
        if (indexValues.size() == 0 ||
                indexValues.size() == 1 && indexValues.get(0) instanceof IntegerConst integerConst && integerConst.getValue() == 0) {
            return pointerValue;
        }
        return buildInst(new GetElementPtrInst(pointerValue, indexValues));
    }

    public Value tryBuildGetElementPtrInst(Value pointerValue, Value... indexValues) {
        if (indexValues.length == 0 ||
                indexValues.length == 1 && indexValues[0] instanceof IntegerConst integerConst && integerConst.getValue() == 0) {
            return pointerValue;
        }
        return buildInst(new GetElementPtrInst(pointerValue, indexValues));
    }

    public BBList<BasicBlock>.BBNode createBasicBlock() {
        BasicBlock basicBlock = new BasicBlock();
        curBasicBlockNodeEnd.insertAfter(basicBlock);
        curBasicBlockNodeEnd = curBasicBlockNodeEnd.getNextNode();
        return curBasicBlockNodeEnd;
    }

    public void switchCurBasicBlock(BBList<BasicBlock>.BBNode basicBlockNode) {
        this.curBasicBlockNode = this.curBasicBlockNodeEnd = basicBlockNode;
        this.endOfBasicBlock = false;
    }

    public String generateLLVM() {
        return moduleIr.parseLLVM();
    }

    public static class ForLoop {
        private final BBList<BasicBlock>.BBNode condBasicBlockNode;
        private final BBList<BasicBlock>.BBNode bodyBasicBlockNode;
        private final BBList<BasicBlock>.BBNode endStmtBasicBlockNode;
        private final BBList<BasicBlock>.BBNode endBasicBlockNode;

        public ForLoop(BBList<BasicBlock>.BBNode condBasicBlockNode, BBList<BasicBlock>.BBNode bodyBasicBlockNode, BBList<BasicBlock>.BBNode endStmtBasicBlockNode, BBList<BasicBlock>.BBNode endBasicBlockNode) {
            this.condBasicBlockNode = condBasicBlockNode;
            this.bodyBasicBlockNode = bodyBasicBlockNode;
            this.endStmtBasicBlockNode = endStmtBasicBlockNode;
            this.endBasicBlockNode = endBasicBlockNode;
        }
    }

    private final Stack<ForLoop> forLoopStack = new Stack<>();
    public void pushForLoop(BBList<BasicBlock>.BBNode condBasicBlockNode, BBList<BasicBlock>.BBNode bodyBasicBlockNode, BBList<BasicBlock>.BBNode endStmtBasicBlockNode, BBList<BasicBlock>.BBNode endBasicBlockNode) {
        forLoopStack.push(new ForLoop(condBasicBlockNode, bodyBasicBlockNode, endStmtBasicBlockNode, endBasicBlockNode));
    }
    public ForLoop popForLoop() {
        return forLoopStack.pop();
    }
    public boolean isInForLoop() {
        return !forLoopStack.isEmpty();
    }
    public BBList<BasicBlock>.BBNode getCondBasicBlockNode() {
        return forLoopStack.peek().condBasicBlockNode;
    }
    public BBList<BasicBlock>.BBNode getBodyBasicBlockNode() {
        return forLoopStack.peek().bodyBasicBlockNode;
    }
    public BBList<BasicBlock>.BBNode getEndStmtBasicBlockNode() {
        return forLoopStack.peek().endStmtBasicBlockNode;
    }
    public BBList<BasicBlock>.BBNode getEndBasicBlockNode() {
        return forLoopStack.peek().endBasicBlockNode;
    }
}
